#!/usr/bin/python3
# -*- coding: utf-8 -*-
"""
@Time        : 2021/11/18 15:57
@Author      : Albert Darren
@Contact     : 2563491540@qq.com
@File        : cluster.py
@Version     : Version 1.0.0
@Description : TODO
@Created By  : PyCharm
"""
from sklearn.cluster import KMeans
from skimage.color import label2rgb
import cv2 as cv
from skimage.color import rgb2gray
from DIP_experiment_5.util import contrast_show

IM_PATH = "../experiment_fig/brain2.tif"
# 读图并转换为灰度图
im = cv.imread(IM_PATH)
img = rgb2gray(im)
# 扁平化并增加一个维度变成列向量
transposed_flattened_im = img.flatten().reshape(-1, 1)
model = KMeans(n_clusters=2, random_state=1)
labels = model.fit_predict(transposed_flattened_im)
labels_im = labels.reshape(img.shape)
rgb_im = label2rgb(labels_im, im, colors=["green"], bg_label=0)
im_dict = {"Origin image": im, "KMeans clustering": rgb_im}
contrast_show(im_dict, (1, 2), mode="brg")
